import torch
import torch.nn as nn
import net
import VoxDataset
from torch import optim
from torch.utils.data import DataLoader
import logging
from torch.autograd import Variable

class Train:
    def __init__(self, mypath):
        #网络结构
        self.in_dim = 3
        self.n_hidden = 16
        self.out_dim = 2
        self.model = net.simpleNet(self.in_dim, self.n_hidden, self.out_dim)
        #超参数
        self.epochs = 48
        self.learnRate = 1e-4
        self.batch_size = 256

        self.trained = False
        self.mypath = mypath

    def training(self):
        #获取数据集
        with open(self.mypath, 'rb') as f:
            self.voxdata = VoxDataset.VoxDataset(f)
        print('ready to train')
        #确认GPU能够运行，否则转移至CPU进行训练
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(device)
        #获取边缘点
        edges = self.voxdata.__getedges__()
        edges = torch.tensor(edges)
        edges.to(device)
        voxedgedata = VoxDataset.VoxedgeDataset(self.voxdata,edges)
        print('边缘点数量')
        print(voxedgedata.__len__())
        print('样本总数量')
        print(self.voxdata.__len__())
        #全局采样数据集
        voxdataloader = DataLoader(dataset=self.voxdata, batch_size=self.batch_size, shuffle=True)
        #边缘采样数据集
        voxedgedataloader = DataLoader(dataset=voxedgedata, batch_size=self.batch_size, shuffle=True)
        #loss函数和优化方式
        loss_func = nn.CrossEntropyLoss()
        loss_func.to(device)
        optimizer = optim.Adam(self.model.parameters(), lr=self.learnRate)

        print('start train')
        for e in range(self.epochs):
            epoch_loss = 0
            self.model.train(True)
            count = 0
            #这里的for循环一次读取的是一个batch的数据,进行随机采样梯度下降
            for batch_idx, (coordinate, flag) in enumerate(voxdataloader):
                if(e > self.epochs / 2 and count > voxdataloader.dataset.__len__() / 4):
                    break
                coordinate, flag = coordinate.to(device), flag.to(device)
                count += self.batch_size

                optimizer.zero_grad()
                flag_pred = self.model(coordinate)
                'print(flag_pred)'
                'print(flag)'
                loss = loss_func(flag_pred, flag)

                loss.backward()
                optimizer.step()

                epoch_loss += loss.item()
            #输出随机采样的loss
            msg = '{}\tEpoch: {}:\t[{}/{}]\tepoch_loss: {:.6f}'.format(
                        'Uniform',
                        e + 1,
                        count,
                        len(self.voxdata),
                        epoch_loss)
            print(msg)
            #当整体结构被模型认知之后，我们的训练重点转为优化细节表示，当训练进度过半，我们进行4分之一的全局随机采样和完整的表面采样
            #表面采样点数取决于模型的表面积，对于简单的模型，转换采样方式将提高训练速度，对于复杂的模型，这种训练方式也可以自适应的提高精度
            if(e > self.epochs / 2):
                for i in range(4):
                    count = 0
                    epoch_loss = 0
                    # 这里专注于表面点的采样
                    for batch_idx, (coordinate, flag) in enumerate(voxedgedataloader):
                        coordinate, flag = coordinate.to(device), flag.to(device)
                        count += self.batch_size

                        optimizer.zero_grad()
                        flag_pred = self.model(coordinate)
                        'print(flag_pred)'
                        'print(flag)'
                        loss = loss_func(flag_pred, flag)
                        loss.backward()
                        optimizer.step()

                        epoch_loss += loss.item()
                    msg = '{}\tEpoch: {}({}/4):\t[{}/{}]\tepoch_loss: {:.6f}'.format(
                        'Surface',
                        e + 1,
                        i + 1,
                        count,
                        len(voxedgedata),
                        epoch_loss)
                    print(msg)
        self.trained = True

    def saving(self, savingpath):
        #参数保存
        torch.save(obj=self.model.state_dict(), f=savingpath)

    def loaing(self, modelpath):
        #从.pth文件中恢复神经网络
        self.model.load_state_dict(torch.load(modelpath))
        #TODO 我应该把它打印出来，并想办法将其画出来。matlab？
        for i in self.model.named_parameters():
            print(i)

    def export(self, num):
        #从神经网络中恢复模型，这里的数据将会被用于模型的重建，没能找到在GPU上运行这一段代码的方式
        voxdata = []
        for z in range(128):
            for y in range(128):
                for x in range(128):
                    temp = torch.FloatTensor([x, y, z])
                    result = self.model(temp)
                    'print(result)'
                    if(result[1] > result[0]):
                        flag = 1
                    else:
                        flag = 0
                    if(len(voxdata) > 0 and flag == voxdata[-2] and voxdata[-1] < 255):
                        voxdata[-1] += 1
                    else:
                        voxdata.append(flag)
                        voxdata.append(1)
        VoxDataset.savingdata(voxdata, self.mypath, num)